#!/usr/bin/python

import pygame
import os.path
import sys

# The resource manager handles requests to load assets and keep them in memory as long as they are needed
# Asset types and their associated return values:
# IMAGE - Surface
# 
class ResourceManager:

    # Wrapper for a controlled resource
    class Resource:
        def __init__(self, data, time):
            self.m_lastUseTime = time
            self.m_data = data
        
        def getLastUseTime(self):
            return self.m_lastUseTime

        def setLastUseTime(self, time):
            self.m_lastUseTime = time

        def getData(self):
            return self.m_data


    # Obvious
    def __init__(self):
        # Init stuff here
        self.m_resourcePaths = set()
        # Loaded resource map is filename_x_y -> { last_use_time, resource data }
        self.m_loadedResources = dict()
        self.m_lastCacheCheck = 0

    # Obvious
    def cleanup(self):
        return None

    # Add a path that will be searched for resources.
    #   Paths will be searched in the order they were added
    #   Paths are used by appending a requested resource filename to each path and attempting load
    def addResourcePath(self, path):
        self.m_resourcePaths.add(path)

    # Per-frame update call - periodically flush unused assets from cache
    def frameUpdate(self, elapsedms, absolutems):

        # Only check once every 15 seconds
        if absolutems < self.m_lastCacheCheck + 15000:
            return

        # Walk each of the resources we know about
        for (name, resource) in self.m_loadedResources.items():
            # Unload if unused within the last 3 minutes
            if absolutems >= resource.getLastUseTime() + 180000:
                # Refcount of 2 means unused - 1 for cache, 1 for this call
                if (sys.getrefcount(resource.getData()) == 2):
                    del self.m_loadedResources[name]

    # Register an image with the resource manager cache
    def registerImage(self, name, surface, defaultDimensions = True):
        dimensions = surface.get_size() 

        # Create resource
        resource = ResourceManager.Resource(surface, pygame.time.get_ticks())

        # Store with dimension mangling
        assetName = "img:"+name+ "_"+str(surface.get_width()) + "_" + str(surface.get_height())
        self.m_loadedResources[assetName] = resource

        # If native resolution, also store without mangling
        if defaultDimensions:
            self.m_loadedResources["img:"+name] = resource


    # Retrieve an image from the resource manager, loading if necessary
    #   returns: Surface
    def loadImage(self, imageDataModule, dimensions = None):
        image = None
        transparencyColor = None
        transparencyThreshold = None
        # Form asset name - base name for image in native size, name_x_y for resized image
        assetName = "img:"+imageDataModule
        if not(dimensions is None):
            assetName = "img:"+imageDataModule+ "_"+str(dimensions[0]) + "_" + str(dimensions[1])

        # Check cache for exact match
        if assetName in self.m_loadedResources:
            # Found
            resource = self.m_loadedResources[assetName]
            resource.setLastUseTime(pygame.time.get_ticks())
            return resource.getData()

        # Fall back and try default size asset in cache
        assetName = "img:"+imageDataModule
        if assetName in self.m_loadedResources:
            image = self.m_loadedResources[assetName].getData()

        # If still not found, try to load
        if image is None:
            # Load the definition module if it exists
            try:
                # Ensure animation data module is loaded
                __import__(imageDataModule)
                module = sys.modules[imageDataModule]

                # Read image data from file
                imageData = getattr(module, 'imageData')
                imageFilename = imageData['image']
                if 'transparencyKeyColor' in imageData:
                    transparencyColor = imageData['transparencyKeyColor']
                if 'transparencyThreshold' in imageData:
                    transparencyThreshold = imageData['transparencyThreshold']
            except ImportError:
                # Probably not a module name, try direct filename
                imageFilename = imageDataModule


            #print "Entity '"+str(assetName)+"' not found in cache. Loading..."
        
            # Try to load from each known path
            for path in self.m_resourcePaths:
                # Form path + filename
                try:
                    fullFilename = os.path.join(path, imageFilename)
                    image = pygame.image.load(fullFilename)
                except pygame.error as exInstance:
                    print "Not able to find %s" % (exInstance)
                    continue # Not in this directory, try next
                
                # Loaded, convert to screen format
                image = image.convert_alpha()
                break

        # Check failed load
        if image is None:
            # Failed load - asset not found
            raise Exception("Could not find asset named '"+imageFilename+"' in any registered resource path.")

        # Convert resolution if needed
        convertNeeded = False
        if not(dimensions is None):
            if (image.get_width() != dimensions[0]) or (image.get_height() != dimensions[1]):
                image = pygame.transform.smoothscale(image, dimensions)
                convertNeeded = True

        # Handle image transparency threshold
        if not(transparencyThreshold is None):
            self.imageThresholdToTransparent(image, transparencyThreshold)

        # If transparency key color is set, set on surface too
        if not(transparencyColor is None):
            image.set_colorkey(transparencyColor)

        # Store in cache
        self.registerImage(imageDataModule, image, (convertNeeded == False))
        
        return image


    # Drop all pixels below a threshold to transparent
    def imageThresholdToTransparent(self, image, threshold):
        for y in range(image.get_height()):
            for x in range(image.get_width()):
                (r,g,b,a) = image.get_at((x,y))
                if r<threshold[0] and g<threshold[1] and b<threshold[2]:
                    image.set_at((x,y), (r,g,b,0))


    # Load frame based animation from a single image based on regular sized frames         
    def loadAnimationFromImage(self, animationDataModule, dimensions = None):
        
        # Form asset name - anim:base name for image in native size, anim:name_x_y for resized image
        assetName = "anim:"+animationDataModule
        if not(dimensions is None):
            assetName = "anim:"+animationDataModule+ "_"+str(dimensions[0]) + "_" + str(dimensions[1])

        # Check cache
        if assetName in self.m_loadedResources:
            # Found
            resource = self.m_loadedResources[assetName]
            resource.setLastUseTime(pygame.time.get_ticks())
            return resource.getData()

        # Ensure animation data module is loaded
        __import__(animationDataModule)
        module = sys.modules[animationDataModule]

        # Read sprite data from file
        spriteData = getattr(module, 'spriteData')
        imageFilename = spriteData['image']
        frameTimems = spriteData['frameTimems']
        sprites = spriteData['sprites']
        
        # Try to load from each known path
        for path in self.m_resourcePaths:
            # Form path + filename
            try:
                fullFilename = os.path.join(path, imageFilename)
                image = pygame.image.load(fullFilename)
            except pygame.error as exInstance:
                continue # Not in this directory, try next
                
            # Loaded, convert to screen format
            image = image.convert_alpha()

            # Slice the image into multiples
            frames = {}
            for spriteName, spriteInfo in sprites.iteritems():
                frames[spriteName] = []
                spriteDimensions = spriteInfo['dimensions']
                coords = spriteInfo['coords']
                for coord in coords:
                    frameSurface = image.subsurface((coord[0], coord[1], spriteDimensions[0], spriteDimensions[1]))
                    frames[spriteName].append(frameSurface)
        
                # Convert resolution of all frames if needed
                convertNeeded = False
                if not(dimensions is None):
                    if (spriteDimensions[0] != dimensions[0]) or (spriteDimensions[1] != dimensions[1]):
                        for frameNum in range(len(frames[spriteName])):
                            frames[spriteName][frameNum] = pygame.transform.smoothscale(frames[spriteName][frameNum], dimensions)
                        convertNeeded = True
            
            # Create resource
            resource = ResourceManager.Resource((frames, frameTimems), pygame.time.get_ticks())

            # Store with dimension mangling
            assetName = "anim:"+imageFilename + "_"+str(image.get_width()) + "_" + str(image.get_height())
            self.m_loadedResources[assetName] = resource

            # If native resolution, also store without mangling
            if convertNeeded == False:
                self.m_loadedResources["anim:"+imageFilename] = resource

            return resource.getData()

        # Failed load - asset not found
        raise Exception("Could not find asset named '"+imageFilename+"' in any registered resource path.")



